import gradio as gr
import torch
from deoldify import device
import warnings
from deoldify.device_id import DeviceId
from deoldify.visualize import *
import tempfile

#choices:  CPU, GPU0...GPU7
device.set(device=DeviceId.GPU0)

plt.style.use('dark_background')
torch.backends.cudnn.benchmark=True

# 忽略警告
warnings.filterwarnings("ignore", category=UserWarning, message=".*?Your .*? set is empty.*?")

# 加载 DeOldify 模型
colorizer = get_image_colorizer(artistic=True)

def colorize_image(input_img, render_factor, watermarked):
    # 创建临时文件
    with tempfile.NamedTemporaryFile(delete=False, suffix='.jpg') as tmp:
        input_img.save(tmp.name)  # 将PIL图像保存到临时文件中
        # 使用临时文件的路径作为输入
        colorized_img = colorizer.get_transformed_image(path=tmp.name, render_factor=render_factor, watermarked=watermarked)
    return colorized_img

# 创建 Gradio 应用
iface = gr.Interface(
    fn=colorize_image,                                        
    inputs=[
        gr.Image(type="pil", label="黑白图像"),                 
        gr.Slider(7, 40, step=1, value=26, label="渲染系数 "),  
    ],
    outputs=gr.Image(type="pil", label="彩色图像"),             
    title="图像上色器",                                         
    description="上传一张黑白照片以通过Deoldify为其上色。",         
    theme="default"                                          
)

# 启动应用
iface.launch()